import torch
import torch.nn as nn

from .basic_modules import CBA


class Stage(nn.Module):
    def __init__(self, inc, ouc, repeat):
        super().__init__()
        self.layer = nn.Sequential(DsBlock(inc, ouc), *[ShuffleBlock(ouc) for _ in range(repeat)])

    def forward(self, x):
        return self.layer(x)


class DsBlock(nn.Module):
    def __init__(self, inc, ouc):
        super().__init__()
        c = ouc // 2
        self.branch1 = nn.Sequential(
            nn.Conv2d(inc, inc, 3, 2, 1, groups=inc, bias=False),
            nn.BatchNorm2d(inc),
            CBA(inc, c, 1, 1),
        )
        self.branch2 = nn.Sequential(
            CBA(inc, c, 1, 1),
            nn.Conv2d(c, c, 3, 2, 1, groups=c, bias=False),
            nn.BatchNorm2d(c),
            CBA(c, c, 1, 1),
        )

    def forward(self, x):
        return torch.cat((self.branch1(x), self.branch2(x)), 1)


class ShuffleBlock(nn.Module):
    def __init__(self, c):
        super().__init__()
        c = c // 2
        self.layers = nn.Sequential(
            CBA(c, c, 1, 1),
            nn.Conv2d(c, c, 3, 1, 1, groups=c, bias=False),
            nn.BatchNorm2d(c),
            CBA(c, c, 1, 1),
        )

    def forward(self, x):
        return torch.cat((x[:, ::2, ...], self.layers(x[:, 1::2, ...])), 1)


class ShuffleNet(nn.Module):
    def __init__(self, channels_list=[16, 32, 64, 128, 256], repeats_list=[1, 2, 3, 1]):
        super().__init__()
        self.cbr = CBA(3, channels_list[0], stride=2)
        self.stages = nn.ModuleList(
            [Stage(*_) for _ in zip(channels_list[:-1], channels_list[1:], repeats_list)]
        )

    def forward(self, x):
        outs = list()
        outs.append(self.cbr(x))
        for stage in self.stages:
            outs.append(stage(outs[-1]))
        return outs[1:]
